#!/usr/bin/env python3
import socket, threading, json, time, random, itertools

HOST = '0.0.0.0'
PORT = 5555
TICK_MS = 500
INITIAL_CASH = 100_000
MAX_LEVERAGE = 10
IMPACT_SIZE = 100
IMPACT_RATIO = 0.01

price = 100.0
users = {}                 # username -> {现金, 仓位列表}
clients = {}
lock = threading.Lock()
next_id = itertools.count()  # 全局自增 id 生成器

# ---------- 工具 ----------
def broadcast(obj):
    msg = (json.dumps(obj) + '\n').encode()
    for sock in list(clients):
        try:
            sock.sendall(msg)
        except:
            disconnect(sock)

def disconnect(sock):
    if sock in clients:
        del clients[sock]
    sock.close()

# ---------- 价格线程 ----------
def price_ticker():
    global price
    while True:
        time.sleep(TICK_MS / 1000)
        price = round(max(0.01, price * (1 + random.uniform(-0.02, 0.02))), 2)
        with lock:
            for u in users.values():
                u['未实现盈亏'] = sum(
                    pos['size'] * (price - pos['price']) for pos in u['仓位列表'])
            broadcast({'type': 'ticker', 'price': price, 'users': len(users)})

# ---------- 主处理 ----------
def handle(sock):
    global price
    username = None
    reader = sock.makefile('r')
    try:
        # 登录
        while True:
            line = reader.readline()
            if not line:
                return
            req = json.loads(line)
            if req['type'] != 'login':
                continue
            username = req['username']
            with lock:
                if username not in users:
                    users[username] = {
                        '现金': INITIAL_CASH,
                        '仓位列表': [],
                        '未实现盈亏': 0
                    }
                clients[sock] = username
            broadcast({'type': 'chat', 'msg': f'{username} 加入了市场'})
            sock.sendall((json.dumps({'type': 'login_ok', 'state': users[username]}) + '\n').encode())
            break

        # 业务循环
        while True:
            req = json.loads(reader.readline())
            u = users[username]

            # 聊天
            if req['type'] == 'chat':
                msg = req['msg'].strip()
                if msg:
                    broadcast({'type': 'chat', 'msg': f'{username}: {msg}'})
                continue

            # 转账
            if req['type'] == 'transfer':
                to, amt = req['to'], int(req['amount'])
                if amt <= 0 or amt > u['现金']:
                    sock.sendall((json.dumps({'type': 'err', 'msg': '转账失败'}) + '\n').encode())
                    continue
                with lock:
                    if to not in users:
                        sock.sendall((json.dumps({'type': 'err', 'msg': '对方不存在'}) + '\n').encode())
                        continue
                    u['现金'] -= amt
                    users[to]['现金'] += amt
                broadcast({'type': 'chat', 'msg': f'{username} 给 {to} 转账 {amt}'})
                continue

            # 开仓
            if req['type'] == 'trade_open':
                side, size, lev = req['side'], int(req['size']), int(req['leverage'])
                if lev < 1 or lev > MAX_LEVERAGE:
                    sock.sendall((json.dumps({'type': 'err', 'msg': '杠杆超限'}) + '\n').encode())
                    continue

                impact_price = price
                if size >= IMPACT_SIZE:
                    sign = 1 if side == 'buy' else -1
                    impact_price = round(price * (1 + sign * IMPACT_RATIO), 2)
                    broadcast({'type': 'chat', 'msg': f'【冲击】{username} 开仓 {size} 股，滑点 {impact_price}'})

                cost = size * impact_price / lev
                if cost > u['现金']:
                    sock.sendall((json.dumps({'type': 'err', 'msg': '资金不足'}) + '\n').encode())
                    continue

                # 记录仓位
                pos = {
                    'id': next(next_id),
                    'side': side,
                    'size': size if side == 'buy' else -size,
                    'price': impact_price,
                    'leverage': lev
                }
                u['现金'] -= cost
                u['仓位列表'].append(pos)

                with lock:
                    price = impact_price
                    for usr in users.values():
                        usr['未实现盈亏'] = sum(
                            p['size'] * (price - p['price']) for p in usr['仓位列表'])
                    broadcast({'type': 'trade_done', 'state': u})

            # 平仓
            if req['type'] == 'trade_close':
                pos_id = int(req['id'])
                pos = None
                for p in u['仓位列表']:
                    if p['id'] == pos_id:
                        pos = p
                        break
                if not pos:
                    sock.sendall((json.dumps({'type': 'err', 'msg': '仓位不存在'}) + '\n').encode())
                    continue

                # 按市价反向成交
                size = abs(pos['size'])
                side = 'sell' if pos['size'] > 0 else 'buy'
                impact_price = price
                if size >= IMPACT_SIZE:
                    sign = 1 if side == 'buy' else -1
                    impact_price = round(price * (1 + sign * IMPACT_RATIO), 2)
                    broadcast({'type': 'chat', 'msg': f'【冲击】{username} 平仓 {size} 股，滑点 {impact_price}'})

                # 结算
                pnl = pos['size'] * (impact_price - pos['price'])
                u['现金'] += pnl
                u['仓位列表'].remove(pos)

                with lock:
                    price = impact_price
                    for usr in users.values():
                        usr['未实现盈亏'] = sum(
                            p['size'] * (price - p['price']) for p in usr['仓位列表'])
                    broadcast({'type': 'trade_done', 'state': u})

    except Exception as e:
        print('client err:', e)
    finally:
        disconnect(sock)
        if username:
            broadcast({'type': 'chat', 'msg': f'{username} 离开了'})

# ---------- 启动 ----------
def main():
    threading.Thread(target=price_ticker, daemon=True).start()
    srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv.bind((HOST, PORT))
    srv.listen()
    print(f'Server listening {HOST}:{PORT}')
    while True:
        conn, addr = srv.accept()
        threading.Thread(target=handle, args=(conn,), daemon=True).start()

if __name__ == '__main__':
    main()
